package main

import (
	"os"
	"go/ast"
	"go/token"
)

func rewrite(funcs []*JunkDrawer) os.Error {
	for _, fn := range funcs {
		c, s := fn.cons, fn.sig

		//package caps
		all := []ast.Stmt{
			rcvr(s.rcvr),
			CaptureAs("t_Ðβご_param_cap_", s.parms),
		}

		//package postconditions
		posts := make([]ast.Stmt, 0, len(c.post)+len(c.in)+1)
		if len(c.pan) > 0 {
			panics := make([]ast.Stmt, 0, len(c.pan))
			for _, pan := range c.pan {
				panics = append(panics, condition(true, s, pan))
			}
			posts = append(posts, wrap_panics(panics))
		}
		for _, in := range c.in {
			posts = append(posts, condition(true, s, in))
		}
		for _, post := range c.post {
			posts = append(posts, condition(true, s, post))
		}
		if len(posts) > 0 {
			all = append(all, defer_wrap(s.rets, posts))
		}

		//package precondtions
		pres := make([]ast.Stmt, 0, len(c.pre)+len(c.in))
		for _, in := range c.in {
			pres = append(pres, condition(false, s, in))
		}
		for _, pre := range c.pre {
			pres = append(pres, condition(false, s, pre))
		}
		if len(pres) > 0 {
			all = append(all, pres...)
		}

		//inject into ast
		fn.fd.Body.List = append(all, fn.fd.Body.List...)
	}
	return nil
}

func rcvr(r string) ast.Stmt {
	if r == "" {
		r = "nil"
	}
	return &ast.DeclStmt{
		&ast.GenDecl{
			Tok: token.VAR,
			Specs: []ast.Spec{
				&ast.ValueSpec{
					Names: []*ast.Ident{
						&ast.Ident{
							Name: "t_Ðβご_rcvr_cap_",
						},
					},
					Type: Empty(),
					Values: []ast.Expr{
						&ast.Ident{
							Name: r,
						},
					},
				},
			},
		},
	}
}

func CaptureAs(name string, nms []string) ast.Stmt {
	var names []ast.Expr
	for _, nm := range nms {
		names = append(names, &ast.Ident{
			Name: nm,
		})
	}
	return &ast.AssignStmt{
		Lhs: []ast.Expr{
			&ast.Ident{
				Name: name,
			},
		},
		Tok: token.DEFINE,
		Rhs: []ast.Expr{
			&ast.CompositeLit{
				Type: &ast.ArrayType{
					Len: nil,
					Elt: Empty(),
				},
				Elts: names,
			},
		},
	}
}

//return the empty interface
func Empty() ast.Expr {
	return &ast.InterfaceType{
		Methods: &ast.FieldList{
			List: []*ast.Field{},
		},
	}
}

//TODO should include type of condition so we can say what kind of contract was
//violated
func condition(dfrd bool, sig *Sig, e Contract) ast.Stmt {
	name := "false"
	rets := &ast.Ident{
		Name: "nil",
	}
	if dfrd {
		name = "true"
		rets.Name = "t_Ðβご_rets_cap_"
	}
	return &ast.IfStmt{
		Cond: &ast.UnaryExpr{
			Op: token.NOT,
			X: &ast.ParenExpr{
				X: e.exp,
			},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.ExprStmt{
					X: &ast.CallExpr{
						Fun: &ast.Ident{
							Name: "f_Ðβご_contract_violation_",
						},
						Args: []ast.Expr{
							&ast.Ident{
								Name: name, //let reporter know to skip
							},
							&ast.BasicLit{
								Kind:  token.STRING,
								Value: e.str,
							},
							&ast.BasicLit{
								Kind:  token.STRING,
								Value: Quote(sig.name),
							},
							&ast.Ident{
								Name: "t_Ðβご_rcvr_cap_",
							},
							&ast.Ident{
								Name: "t_Ðβご_param_cap_",
							},
							rets,
						},
					},
				},
			},
		},
	}
}

//take all of the postconditional, panicconditional, and invariantconditional
//statements and wrap them up in a defer'd func to insert at the head of the
//func body to be annotated
func defer_wrap(caps []string, posts []ast.Stmt) ast.Stmt {
	//defer func() {
	//  rets := []interface{}{ret,ar,gs}
	//	posts
	//}()
	posts = append([]ast.Stmt{CaptureAs("t_Ðβご_rets_cap_", caps)}, posts...)
	return &ast.DeferStmt{
		Call: &ast.CallExpr{
			Fun: &ast.FuncLit{
				Type: &ast.FuncType{
					Params: &ast.FieldList{}, //this is required
				},
				Body: &ast.BlockStmt{
					List: posts,
				},
			},
		},
	}
}

func wrap_panics(panics []ast.Stmt) ast.Stmt {
	// gist, though we use something far less likely than x
	// if x := recover(); x != nil {
	//    posts
	// }
	return &ast.IfStmt{
		// x := recover()
		Init: &ast.AssignStmt{
			Lhs: []ast.Expr{
				ast.NewIdent("_Ðβご_recover"),
			},
			Tok: token.DEFINE,
			Rhs: []ast.Expr{
				&ast.CallExpr{
					Fun: &ast.Ident{
						Name: "recover",
					},
				},
			},
		},
		// x != nil
		Cond: &ast.BinaryExpr{
			X: &ast.Ident{
				Name: "_Ðβご_recover",
			},
			Op: token.NEQ,
			Y: &ast.Ident{
				Name: "nil",
			},
		},
		Body: &ast.BlockStmt{
			List: append(panics, &ast.ExprStmt{
				X: &ast.CallExpr{
					Fun: &ast.Ident{
						Name: "panic",
					},
					Args: []ast.Expr{
						&ast.Ident{
							Name: "_Ðβご_recover",
						},
					},
				},
			}),
		},
	}
}
